﻿using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Text;
using System.Runtime.InteropServices;
using System.Collections;
using MPI;

namespace Dwarf.Managed
{
    /// <summary>
    /// User interface for solving the Unstructured Grid Problem.
    /// </summary>
    class Program
    {
        /// <summary>
        /// The settings for this dwarf.
        /// </summary>
        static private Settings settings;

        /// <summary>
        /// Send an array of type T to target
        /// </summary>
        /// <typeparam name="T">
        /// Element type of array
        /// </typeparam>
        /// <param name="values">
        /// array to send</param>
        /// <param name="target">
        /// who to send to
        /// </param>
        static public void SendArray<T>(T[] values, int target)
        {
            //length as tag
            Communicator.world.Send(values, target, values.Length);
        }

        /// <summary>
        /// Return an array of type T[] received from source
        /// </summary>
        /// <typeparam name="T">
        /// Element type of array
        /// </typeparam>
        /// <param name="source">
        /// Who to receive array from
        /// </param>
        /// <param name="status">
        /// Completed status of Receive operation
        /// </param>
        /// <returns>
        /// A newly allocated array containing the data received
        /// </returns>
        static public T[] ReceiveArray<T>(int source, out MPI.CompletedStatus status)
        {
            //probe for length and then allocate array to receive
            Status probe = Communicator.world.Probe(source, Communicator.anyTag);
            T[] result = new T[probe.Tag];
            Communicator.world.Receive(probe.Source, probe.Tag, ref result, out status);
            return result;
        }

        /// <summary>
        /// Returns an array of type T[] received from source
        /// </summary>
        /// <typeparam name="T">
        /// Element type of received array
        /// </typeparam>
        /// <param name="source">
        /// Who to receive it from
        /// </param>
        /// <returns>
        /// A newly allocated array containing the data received
        /// </returns>
        static public T[] ReceiveArray<T>(int source)
        {
            MPI.CompletedStatus status;
            return ReceiveArray<T>(source, out status);
        }

        /// <summary>
        /// Init & fill neigbourhood array
        /// </summary>
        static private void constructNeighborhood(Solver ugsolver)
        {
            //Init map
            HashSet<int>[] neighbors = new HashSet<int>[ugsolver.numVertex];
            for (int j = 0; j < ugsolver.numVertex; j++)
            {
                neighbors[j] = new HashSet<int>();
            }

            //Insert verteces into map
            for (int j = 0; j < ugsolver.numTriangles; j++)
            {
                neighbors[ugsolver.triangles[0][j]].Add(ugsolver.triangles[1][j]);
                neighbors[ugsolver.triangles[0][j]].Add(ugsolver.triangles[2][j]);

                neighbors[ugsolver.triangles[1][j]].Add(ugsolver.triangles[0][j]);
                neighbors[ugsolver.triangles[1][j]].Add(ugsolver.triangles[2][j]);

                neighbors[ugsolver.triangles[2][j]].Add(ugsolver.triangles[0][j]);
                neighbors[ugsolver.triangles[2][j]].Add(ugsolver.triangles[1][j]);
            }

            //Convert from hashed maps to arrays
            ugsolver.neighborhood = neighbors.Select((n) => n.ToArray()).ToArray();
        }

        /// <summary>
        /// Construct array of shadow (boundary) vertices for specified proc.
        /// </summary>
        /// <param name="ugsolver">
        /// Instance of Solver.
        /// </param>
        /// <param name="vertices">
        /// Aray for each process listing vertices in that process.
        /// </param>
        /// <param name="processMap">
        /// Array mapping vertex to process containing vertex
        /// </param>
        /// <param name="procNum">
        /// MPI process for which to construct maps.
        /// </param>
        static private int[][] constructShadowMaps(
            Solver ugsolver,
            int[][] vertices,
            int[] processMap,
            int procNum)
        {
            //create set of all neighbors of all vertices in proc
            HashSet<int> vSet = new HashSet<int>();
            foreach (int v in vertices[procNum])
            {
                vSet.UnionWith(ugsolver.neighborhood[v]);
            }

            //Remove vertices in this proc
            //result is neighbors in other procs, the ones we need to shadow
            vSet.ExceptWith(vertices[procNum]);

            //group by proc that contains shadow giving [proc][shadows we need from proc]
            var procShadows = new HashSet<int>[vertices.Length].Select(junk => new HashSet<int>()).ToArray();
            foreach (int v in vSet) { procShadows[processMap[v]].Add(v); }
            return procShadows.Select((shadows) => shadows.ToArray()).ToArray();
        }

        /// <summary>
        /// Split graph of vertices greedily, i.e grab neighbors if possible
        /// </summary>
        /// <param name="ugsolver">
        /// Instance of Solver.
        /// </param>
        /// <param name="vertices">
        /// Vertices for each process.
        /// </param>
        /// <returns>
        /// Map of vertex to process
        /// </returns>
        static private int[] greedyGraphDivision(Solver ugsolver, int[][] vertices)
        {
            //allocate array to return so we can fill it
            int[] processMap = new int[ugsolver.numVertex];

            //get an enumerator of un-allocated vertices
            int procs = vertices.Length;
            int pendingFlag = -procs;
            int not_assigned = pendingFlag - 1;
            processMap = processMap.Select((junk) => not_assigned).ToArray();  //fill with -2
            IEnumerator<int> unassigned = new MapIterator(processMap).GetEnumerator();
                     
            //Compute number of vertices for each process
            //and fill it
            for (int proc = 0; proc < procs; proc++)
            {
                //compute number of vertices for proc
                int num = ugsolver.numVertex;
                int n = num / procs + ((num % procs > proc) ? 1 : 0);
                vertices[proc] = new int[n];

                //breadth first neighbor walk (using queue)
                Queue<int> neighbors = new Queue<int>();
                int i = 0;
                while (i < vertices[proc].Length)
                {
                    int next = -1;
                    //take from queue if possible
                    if (neighbors.Count() > 0)
                    {
                        //if already assigned, try again
                        next = neighbors.Dequeue();
                        if (processMap[next] >= 0) continue;
                    }
                    else
                    {
                        if (!unassigned.MoveNext()) throw new Exception("improper division, ran out of vertices");
                        next = unassigned.Current;
                    }

                    //assign to this proc
                    vertices[proc][i] = next;
                    processMap[next] = proc;
                    i++;

                    //add neighbors to queue if not assigned or pending yet
                    foreach (int v in ugsolver.neighborhood[next])
                    {
                        if (processMap[v] < pendingFlag)
                        {
                            neighbors.Enqueue(v);
                            processMap[v] = pendingFlag;
                        }
                    }
                }
                pendingFlag++;  //treat prior pending as unassigned
            }
            int sanity = vertices.Select((verts) => verts.Length).Sum();
            if (sanity != ugsolver.numVertex) throw new Exception("improper division of graph");
            
            return processMap;
        }


        /// <summary>
        /// Split graph & distribute parts to processes.
        /// </summary>
        /// <param name="ugsolver">
        /// Instance of Solver.
        /// </param>
        static private void splitGraph(Solver ugsolver)
        {
            if (Communicator.world.Rank == Settings.ROOT_PROCESS)
            {
             
                int[][] vertices = new int[Communicator.world.Size][];

                //Greedy graph division
                int[] processes = greedyGraphDivision(ugsolver, vertices);

                //destribution sub-graphs to processes
                for (int proc = 1; proc < Communicator.world.Size; proc++)
                {
                    //send vertices to process
                    SendArray(vertices[proc], proc);

                    //collect values for vertices and send to process
                    double[] tempValues = new double[vertices[proc].Length];
                    for (int j = 0; j < vertices[proc].Length; j++)
                    { tempValues[j] = ugsolver.vertexValues[vertices[proc][j]]; }                    
                    SendArray(tempValues, proc);

                    //send neighborhood for vertices to process
                    for (int j = 0; j < vertices[proc].Length; j++)
                    {
                        int v = vertices[proc][j];
                        SendArray(ugsolver.neighborhood[v], proc);
                    }            

                    //numbers of shadow vertices are being sended to process 
                    int[][] shadows = constructShadowMaps(ugsolver, vertices, processes, proc);
                    for (int j = 0; j < shadows.Length; j++)
                    { 
                        SendArray(shadows[j], proc);
                    }
                }    

                //computation of root data

                //values of vertices for root process
                ugsolver.localVertices = vertices[0];
                ugsolver.localValues = new Dictionary<int, double>(ugsolver.localVertices.Length);
                for (int j = 0; j < ugsolver.localVertices.Length; j++)
                { 
                    int v = ugsolver.localVertices[j];
                    ugsolver.localValues[v] = ugsolver.vertexValues[v]; 
                }

                //copy neighbor's numbers into localNeighborhood
                ugsolver.localNeighborhood = new Dictionary<int,int[]>(ugsolver.localVertices.Length);
                for (int j = 0; j < ugsolver.localVertices.Length; j++)
                {
                    int v = ugsolver.localVertices[j];
                    ugsolver.localNeighborhood[v] = ugsolver.neighborhood[v];
                }

                //numbers of shadow (boundary) vertices for root
                ugsolver.localShadows = constructShadowMaps(ugsolver, vertices, processes, 0);
            } 
            else
            {
                //receive local vertices
                ugsolver.localVertices = ReceiveArray<int>(Communicator.anySource);

                //receive values of local vertices
                double[] tempValues = ReceiveArray<double>(Communicator.anySource);

                //fill local values and neighbors map
                ugsolver.localValues = new Dictionary<int, double>(ugsolver.localVertices.Length);
                ugsolver.localNeighborhood = new Dictionary<int,int[]>(ugsolver.localVertices.Length);
                for (int j = 0; j < ugsolver.localVertices.Length; j++)
                {
                    int v = ugsolver.localVertices[j];
                    ugsolver.localValues[v] = tempValues[j];
                    ugsolver.localNeighborhood[v] = ReceiveArray<int>(Communicator.anySource);
                }

                //Receive shadow vertices maps
                ugsolver.localShadows = new int[Communicator.world.Size][];
                for (int j = 0; j < ugsolver.localShadows.Length; j++)
                { 
                    ugsolver.localShadows[j] = ReceiveArray<int>(Communicator.anySource);
                }
            }
            ugsolver.localize();
        }

        /// <summary>
        /// Get settings from the command line.
        /// </summary>
        /// <param name="args">
        /// Cli params.
        /// </param>
        /// <param name="ugsolver">
        /// Instance of Solver.
        /// </param>
        static private void getSettings(String[] args, Solver ugsolver)
        {
            bool error = false;
            if (Communicator.world.Rank == Settings.ROOT_PROCESS)
            {
                try
                {
                    settings = new Settings();
                    settings.Init(args, ugsolver);      //Parse all arguments for dwarf.

                    constructNeighborhood(ugsolver);
                }
                catch (Exception ex)
                {
                    Console.WriteLine(ex.Message);
                    error = true;
                }
            }

            //Distribute error if that occurred while parsing.
            Communicator.world.Broadcast(ref error, Settings.ROOT_PROCESS);
            if (error) throw new Exception("");

            //Distribute epsilon.
            Communicator.world.Broadcast(ref ugsolver.epsilon, Settings.ROOT_PROCESS);

            //split and distribute graph to MPI processes
            splitGraph(ugsolver);

            //drop no longer needed data
            ugsolver.neighborhood = null;
            ugsolver.triangles = null;
        }

        /// <summary>
        /// Point of the program start.
        /// </summary>
        /// <param name="args">
        /// Cli params.
        /// </param>
        static void Main(String[] args)
        {
            try
            {
                using (new MPI.Environment(ref args))                       // MPI Initialization.
                {
                    if (Communicator.world.Size == 1)
                    {
                        throw new Exception("Only one mpi process.");
                    }

                    Solver ugsolver = new Solver();

                    getSettings(args, ugsolver);                            // Get settings from the command line.

                    Communicator.world.Barrier();
                    if (Communicator.world.Rank == Settings.ROOT_PROCESS)
                        settings.Start();                                   // Start new time count.
                    ugsolver.Solve();                                       // Solve the current problem.
                    Communicator.world.Barrier();
                    ugsolver.CollectResults();
                    if (Communicator.world.Rank == Settings.ROOT_PROCESS)
                        settings.Finish(ugsolver);                          // Stop the time count and write results.
                }

            }
            catch (Exception ex)
            {
                Console.WriteLine(ex.Message);
                Console.WriteLine(ex.StackTrace);
            }
        }
    }


    /// <summary>
    /// An iterator across the processMap array looking for unallocated vertices
    /// (negative values). At each step, it yields the the index within the processMap
    /// where a negative (unallocated) value is stored.
    /// </summary>
    class MapIterator
    {
        private int[] processMap;
        private int current;
        public MapIterator(int[] processMap)
        { this.processMap = processMap; }

        public System.Collections.Generic.IEnumerator<int> GetEnumerator()
        {
            for (current = 0; current < processMap.Length; current++)
            {
                if (processMap[current] < 0) yield return current;
            }
        }
    }
}